package com.witmore.neutron.ai.cognitive;

import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.http.codec.ServerSentEvent;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Description: ChatFluxUtil
 * Author: 豆子高
 * Date: 2025/8/20 18:04
 */
@Slf4j
public class ChatFluxUtilSSe {


    public static Flux<ServerSentEvent<String>> returnChatFlux(Flux<ChatResponse> chatResponseFlux, ChatMemory chatMemory, String sessionId) {
        // 用于保存最后一个响应块
        AtomicReference<ChatResponse> lastResponse = new AtomicReference<>();
        AtomicReference<AssistantMessage> assistantMessage = new AtomicReference<>();
        return Flux.create((FluxSink<ServerSentEvent<String>> emitter) -> {
            long pipelineStart = System.currentTimeMillis();
            chatResponseFlux
                    .doOnSubscribe(s -> log.info("pipeline started"))
                    .doOnNext(response -> {
                        // 保存最后一个响应块
                        lastResponse.set(response);

                        // 原有的内容发送逻辑
                        if (response == null || response.getResult() == null) {
                            emitter.complete();
                            return;
                        }
                        AssistantMessage output = response.getResult().getOutput();
                        emitter.next(buildMessageEvent(
                                output.getMessageType().getValue(),
                                output.getText() == null ? "" : output.getText(),
                                response.getResult().getMetadata().getFinishReason()
                        ));
                        assistantMessage.set(output);
                    })
                    .doOnComplete(() -> {
                        // 流结束时，检查最后一个响应块中是否有usage信息
                        ChatResponse finalResponse = lastResponse.get();
                        if (finalResponse != null) {
                            Usage finalUsage = finalResponse.getMetadata().getUsage();
                            if (finalUsage != null) {
                                log.info("Final Usage Stream: {}", finalUsage);
                                Map<String, Object> metadata = assistantMessage.get().getMetadata();
                                metadata.put("usage", finalUsage);
                                chatMemory.add(sessionId, assistantMessage.get());
                            }
                        }
                        emitter.next(buildEndEvent());
                        emitter.complete();
                        log.info("Pipeline completed in {} ms", System.currentTimeMillis() - pipelineStart);
                    })
                    .doOnError(error -> {
                        emitter.next(buildErrorEvent(error.getMessage()));
                        emitter.complete();
                    })
                    .subscribe(null, error -> {
                        log.error("流处理异常", error);
                        emitter.next(buildErrorEvent(error.getMessage()));
//                      emitter.complete();
                    });
        }, FluxSink.OverflowStrategy.BUFFER);


    }

    // 辅助方法封装
    public static ServerSentEvent<String> buildMessageEvent(String role, String content, String reason) {
        return ServerSentEvent.<String>builder()
                .data(JSONObject.toJSONString(Map.of(
                        "role", role,
                        "content", content,
                        "finishReason", reason
                )))
                .event("message")
                .build();
    }

    public static ServerSentEvent<String> buildErrorEvent(String msg) {
        return ServerSentEvent.<String>builder()
                .data("{\"error\":\"" + msg + "\"}")
                .event("error")
                .build();
    }

    public static ServerSentEvent<String> buildEndEvent() {
        return ServerSentEvent.<String>builder()
                .event("end")
                .data("{\"status\":\"completed\"}")
                .build();
    }
}
